using Test
using LinearAlgebra
using FLEXINVERT.PriorCovariance
using FLEXINVERT.Transformations

const DLM = Base.require(Base.PkgId(Base.UUID("8bb1440f-4735-579b-a4ab-409b98df4dab"), "DelimitedFiles"))

function load_vector_with_optional_index(path::String)
    data = DLM.readdlm(path)
    if ndims(data) == 2 && size(data, 2) == 2 && all(data[:, 1] .== collect(1:size(data, 1)))
        return vec(data[:, 2])
    end
    return vec(data)
end

@testset "Fortran parity for transforms" begin
    root = normpath(joinpath(@__DIR__, "..", "..", ".."))
    fortran_dir = joinpath(root, "run", "work", "minimal_env", "flexinvert_run", "FIVOUT_05deg")
    debug_dir = joinpath(root, "run", "work", "minimal_env", "flexinvert_run", "FIVOUT_05deg_debug")

    @test isdir(fortran_dir)
    @test isdir(debug_dir)

    eigenvalues = vec(DLM.readdlm(joinpath(fortran_dir, "evals.txt")))
    eigenvectors = DLM.readdlm(joinpath(fortran_dir, "evecs.txt"))
    sqcort_raw = DLM.readdlm(joinpath(fortran_dir, "sqcort.txt"))
    cort_raw = DLM.readdlm(joinpath(fortran_dir, "cort.txt"))
    icort_raw = DLM.readdlm(joinpath(fortran_dir, "icort.txt"))

    ntstate = size(sqcort_raw, 1)
    n_spatial = size(eigenvectors, 1)
    nmodes = length(eigenvalues)

    temporal_eig = eigen(Symmetric(cort_raw))
    sqcort = copy(sqcort_raw)
    isqcort = temporal_eig.vectors * Diagonal(1.0 ./ sqrt.(temporal_eig.values)) * temporal_eig.vectors'

    sqrt_eigs = sqrt.(eigenvalues)
    inv_sqrt_eigs = 1.0 ./ sqrt_eigs
    inv_eigs = 1.0 ./ eigenvalues

    cov = CovarianceMatrix(
        nothing,
        eigenvalues,
        eigenvectors,
        cort_raw,
        icort_raw,
        sqcort,
        isqcort,
        nmodes,
        sum(eigenvalues),
        1.0,
        n_spatial,
        1,
        ntstate,
        n_spatial,
        ones(Float64, n_spatial * ntstate),
        ones(Float64, n_spatial * ntstate),
        false
    )

    phi_vec = load_vector_with_optional_index(joinpath(fortran_dir, "phi_solution.txt"))
    chi_vec = load_vector_with_optional_index(joinpath(debug_dir, "chi_solution.txt"))

    @test length(phi_vec) == length(chi_vec)

    chi_calc = Transformations.phi2chi(phi_vec, cov)
    @test chi_calc ≈ chi_vec rtol=1e-12 atol=1e-12

    phi_back = Transformations.chi2phi(chi_vec, cov)
    @test phi_back ≈ phi_vec rtol=1e-12 atol=1e-12

    spatial_bhalf = eigenvectors * Diagonal(inv_sqrt_eigs) * eigenvectors'
    B_half_matrix = kron(isqcort, spatial_bhalf)
    chi_linear = B_half_matrix * phi_vec
    @test chi_linear ≈ chi_vec rtol=1e-12 atol=1e-12

    precision_spatial = eigenvectors * Diagonal(inv_eigs) * eigenvectors'
    precision_matrix = kron(icort_raw, precision_spatial)
    binv_calc = Transformations.apply_B_inv(chi_vec, cov)
    precision_expected = precision_matrix * chi_vec
    @test binv_calc ≈ precision_expected rtol=1e-12 atol=1e-12

    phys_len = length(chi_vec)
    n_ctrl = length(phi_vec)
    B_half_loop = Matrix{Float64}(undef, phys_len, n_ctrl)
    basis = zeros(Float64, n_ctrl)
    column = zeros(Float64, phys_len)
    for j in 1:n_ctrl
        basis .= 0.0
        basis[j] = 1.0
        Transformations.phi2chi!(column, basis, cov)
        @inbounds B_half_loop[:, j] = column
    end
    @test B_half_loop ≈ B_half_matrix rtol=1e-12 atol=1e-12

    hmat = readdlm(joinpath(fortran_dir, "hmat.txt"))
    H_reduced_calc = Transformations.build_H_reduced(hmat, cov)
    H_reduced_expected = hmat * B_half_loop
    @test H_reduced_calc ≈ H_reduced_expected rtol=1e-12 atol=1e-12
end
